package main

import (
	"flag"
	"fmt"
	"github.com/gorilla/websocket"
	"github.com/jiangluu/log"
	"github.com/jiangluu/log/handlers/cli"
	"io/ioutil"
	logOld "log"
	"strings"
	"sync"
	"sync/atomic"
	"time"
)

const (
	message_file        = "./message.txt"
	interval            = 1
	connectionTimeout   = 10
	template_accountKey = "##ACCKEY##"
)

var (
	duration    = 600
	arr_message []string
	log1        *log.Entry
	rand1       = 1
)

func init() {
	aa, err := ioutil.ReadFile(message_file)
	if nil != err {
		panic(err)
	}

	rand1 = time.Now().Nanosecond() % 1000

	ss := string(aa)
	a_ss := strings.Split(ss, "\r\n")
	arr_message = a_ss

	log.SetHandler(cli.Default)
	log.SetLevel(log.InfoLevel)
}

func main() {
	al := log.WithField("action", "main")
	log1 = al

	al.Info("init")

	addr := flag.String("a", "wss://kun0.wdfunny.com/", "endpoint of server")
	concurent := flag.Int("c", 5, "concurent workers")
	f3 := flag.Int("d", 30, "duration")
	flag.Parse()

	duration = *f3

	al.WithField("endpoint", *addr).WithField("concurent", *concurent).WithField("duration", duration).Info("flag read")

	WsBench(*addr, *concurent)
}

func getMyMess(ind int) []string {
	N := len(arr_message)
	rr := make([]string, N)
	copy(rr, arr_message)

	bb := fmt.Sprintf("bot%03d%04d", rand1, ind)

	for i := 0; i < N; i++ {
		rr[i] = strings.Replace(rr[i], template_accountKey, bb, -1)
	}

	return rr
}

func WsBench(address string, sockets int) {
	log1.WithField("endpoint", address).Info("connecting")

	start := time.Now()
	counter := 0
	var connectionSuc int32 = 0
	var readCounter uint64 = 0
	var writeCounter uint64 = 0
	var readError uint64 = 0
	var writeError uint64 = 0
	var connectionError uint64 = 0
	var writeBytes uint64 = 0
	var readBytes uint64 = 0
	var durr time.Duration
	var wg sync.WaitGroup

	for safer := 0; safer < sockets; safer++ {
		counter++
		wg.Add(1)

		go func(my_i int) {
			defer wg.Done()
		lable1:
			dialer := websocket.Dialer{
				HandshakeTimeout: time.Duration(connectionTimeout) * time.Second,
			}
			co, _, err := dialer.Dial(address, nil)

			if err != nil {
				log1.WithError(err).Error("Dial")
				connectionError++
				goto lable1
			} else {
				atomic.AddInt32(&connectionSuc, 1)
				defer co.WriteMessage(websocket.CloseMessage, websocket.FormatCloseMessage(websocket.CloseNormalClosure, ""))

				//connectionSucF := atomic.LoadUint64(&connectionSuc)
				//log1.WithField("connectionSuccess", connectionSucF).Info("connectionSuc")

				my_mess := getMyMess(my_i)
				num := len(my_mess)
				count2 := 0
				for { // LOOP
					if time.Since(start) > (time.Duration(duration) * time.Second) {
						break
					}
					message := my_mess[num-1]
					if count2 < num-1 {
						message = my_mess[count2]
					}
					log1.WithField("M", message).WithField("c", count2).WithField("I", my_i).Debug("mymess")
					count2 += 1

					writeTime := time.Now()
					err = co.WriteMessage(websocket.TextMessage, []byte(message))
					if err != nil {
						log1.WithError(err).Error("write")
						atomic.AddUint64(&writeError, 1)
					} else {
						atomic.AddUint64(&writeBytes, uint64(len([]byte(message))))
						atomic.AddUint64(&writeCounter, 1)
						_, readMessage, err := co.ReadMessage()
						if err != nil {
							log1.WithError(err).Error("read")
							readError++
						} else {
							log1.WithField("A", string(readMessage)).Debug("Ack")
							atomic.AddUint64(&readBytes, uint64(len([]byte(readMessage))))
							atomic.AddUint64(&readCounter, 1)
						}
					}
					dur := time.Since(writeTime)
					durr += dur
					time.Sleep(time.Duration(interval) * time.Second)
				}
			}
		}(safer + 1)

	}

	go func() {
		for {
			readCounterF := atomic.LoadUint64(&readCounter)
			writeCounterF := atomic.LoadUint64(&writeCounter)
			writeBytesF := atomic.LoadUint64(&writeBytes)
			readBytesF := atomic.LoadUint64(&readBytes)
			connectionErrorF := atomic.LoadUint64(&connectionError)
			writeErrorF := atomic.LoadUint64(&writeError)
			readErrorF := atomic.LoadUint64(&readError)
			connectionSucF := atomic.LoadInt32(&connectionSuc)
			var averageRtt time.Duration
			if readCounterF == 0 {
				averageRtt = time.Duration(0)
			} else {
				averageRtt = durr / time.Duration(readCounterF)
			}
			logOld.Println("connectionSuc:", connectionSucF, ",Total Sent:", writeCounterF, ", Total Received:", readCounterF, ", Bytes Sent", writeBytesF, ", Bytes Received:", readBytesF, ", Average RTT:", averageRtt, ", Connection Error:", connectionErrorF, ", Write Error:", writeErrorF, ", Read Error:", readErrorF)

			time.Sleep(time.Duration(1) * time.Second)
		}
	}()

	wg.Wait()
}
